# 根据量减价升的原则进行交易
# 当前价格处于均线之下时，上涨概率应该会更高一点
import talib
import pandas as pd
from datacache import get_stock_data
import matplotlib.pyplot as plt
import mplfinance as mpf

def plot_results(df):
    mpf_style = mpf.make_mpf_style(base_mpf_style='charles', rc={ 'font.family': 'SimHei', 'axes.unicode_minus': 'False' })
    df = df.rename(
        {
            "date": "Date",
            "open": "Open",
            "close": "Close",
            "low": "Low",
            "high": "High",
            "volume": "Volume",
        },
        axis=1,
    )
    # 将 Date 列设为索引
    df.index = df["Date"].astype("datetime64[ns]")
    df = df.sort_index()
    # 将 Date 列设为索引
    # 创建子图
    fig, (ax1, ax2, ax3, ax4) = plt.subplots(4, 1, figsize=(30, 25))

    # 绘制 K 线图和均线
    ap = [
        mpf.make_addplot(df['ma5'], color='blue', width=0.7, ax=ax1),
        mpf.make_addplot(df['ma10'], color='orange', width=0.7, ax=ax1),
        mpf.make_addplot(df['ma20'], color='green', width=0.7, ax=ax1),
        mpf.make_addplot(df['Volume'] / 100000000, color='green', width=0.7, ax=ax3),

        mpf.make_addplot(df['macd'], panel=1, color='#2d90fd', title='MACD', type="bar", ax=ax2),
        mpf.make_addplot(df['macd_signal'], panel=1, color='red', ax=ax2),
        # mpf.make_addplot(df['macd_hist'], panel=1, type='bar', color='dimgray', ax=ax3)
    ]
    mpf.plot(df, type='candle', style=mpf_style, addplot=ap, volume=ax3, ax=ax1)
    ax1.set_title('K Line')

    # 绘制策略收益图
    ax4.plot(df.index, df['sum'], label='total profit rate', color='purple')
    # ax4.axhline(0, color='gray', linestyle='--', linewidth=0.5)
    ax4.set_title('strategy profit')
    ax4.legend()

    # 调整布局
    plt.tight_layout()
    plt.show()

def main():
    stock_code = 'sh600116'  # 平安银行
    start_date = '20230801'
    end_date = '20240801'

    df = get_stock_data(stock_code, start_date, end_date)
    macd, signal, hist = talib.MACD(df['close'], fastperiod=12, slowperiod=26, signalperiod=9)
    macd_series = pd.Series(macd)
    macd_filled = macd_series.fillna(0).values
    macd_signal_series = pd.Series(macd)
    macd_signal_filled = macd_signal_series.fillna(0).values
    df['sum'] = 0
    df['macd'] = macd_filled
    df['macd_signal'] = macd_signal_filled
    df['macd_hist'] = hist
    df['mean5'] = df['volume'].rolling(window=5).mean() # 前5日平均成交量
    initial_capital = 10000  # 初始资金
    holdings = 0
    holding_price = 0 # 持仓价格
    for index, row in df.iterrows():
        close_price = row['close']
        open_price = row['open']
        high_price = row['high']  # 最高价
        low_price = row['low']  # 最低价
        volume = row['volume']
        turnover = row['turnover']  # 换手率
        ma10 = row['ma10']  # 10日均线
        ma5 = row['ma5']  # 5日均线
        mean5 = row['mean5']  # 前5日平均成交量
        today_macd = row['macd']
        if index > 0:
            last_day_macd = df.loc[index - 1, 'macd']
            if holdings > 0:  # 如果有持仓 准备卖出
                if today_macd < last_day_macd: # 如果今天macd小于0，昨天macd大于0，卖出
                    initial_capital = initial_capital + holdings * close_price
                    holding_price = 0
                    holdings = 0
            else: # 如果没有持有股票，准备买入
                if today_macd > 0 and last_day_macd <= 0: # 如果今天macd大于0，昨天macd小于0，买入
                    holdings = round(initial_capital / close_price, 0)
                    initial_capital = initial_capital - holdings * close_price
                    holding_price = close_price
                pass
        df.loc[index, 'sum'] = initial_capital + holdings * close_price
    plot_results(df)

if __name__ == "__main__":
    main()